import '../lib/utils/loadEnv';
import { StateGraph, END, Annotation, START } from '@langchain/langgraph';
import { HumanMessage, AIMessage, BaseMessage } from '@langchain/core/messages';
import { getOpenaiInstance } from '../lib/llm/openai';

console.log('%c Line:6 🥕', 'color:#f5ce50', process.env.OPENAI_API_KEY);
const StateAnnotation = Annotation.Root({
  messages: Annotation<BaseMessage[]>({
    reducer: (state, update) => {
      return state.concat(update);
    },
  }),

  userInput: Annotation<string>(),
});

const app = () => {
  const llm = getOpenaiInstance();
  // 创建状态图
  const graph = new StateGraph(StateAnnotation)
    .addNode('processUserInput', async (state) => {
      // 创建一个新的人类消息
      const humanMessage = new HumanMessage(state.userInput);
      // 将人类消息添加到历史中
      console.log(
        '%c Line:24 🍑 humanMessage，收到消息，正在处理',
        'color:#b03734',
        humanMessage
      );
      // 返回更新后的消息历史
      return { messages: [humanMessage] };
    })
    .addNode('generateAIResponse', async (state) => {
      // 使用 LLM 生成响应
      const response = await llm.invoke(state.messages);

      // 将 AI 响应添加到历史中
      const newMessages = [...state.messages, response];

      // 返回更新后的消息历史
      return { messages: newMessages };
    })
    .addEdge(START, 'processUserInput')
    .addEdge('processUserInput', 'generateAIResponse')
    .addEdge('generateAIResponse', END);

  // 编译图
  return graph.compile();
};

// 使用函数运行对话
async function runConversation(userInput) {
  // 调用图，传入用户输入
  const result = await app().invoke({ userInput });
  console.log('%c Line:55 🍩 result', 'color:#2eafb0', result);

  // 获取最后一条消息（AI 的响应）
  const lastMessage = result.messages[result.messages.length - 1];

  // 返回 AI 的响应内容
  return lastMessage.content;
}

// 示例使用
async function main() {
  // 第一轮对话
  const response1 = await runConversation('你好，请介绍一下自己');
  console.log('AI 响应:', response1);

  // 第二轮对话（注意：这将是一个新的对话，因为我们没有保持状态）
  const response2 = await runConversation('什么是 LangGraphJS？');
  console.log('\n\n\nAI 响应:', response2);
}

main().catch(console.error);
